/*---------------------------------------------------------------------------*\
  =========                 |
  \\      /  F ield         | OpenFOAM: The Open Source CFD Toolbox
   \\    /   O peration     | Website:  https://openfoam.org
    \\  /    A nd           | Copyright (C) 2025 OpenFOAM Foundation
     \\/     M anipulation  |
-------------------------------------------------------------------------------
License
    This file is part of OpenFOAM.

    OpenFOAM is free software: you can redistribute it and/or modify it
    under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    OpenFOAM is distributed in the hope that it will be useful, but WITHOUT
    ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
    FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
    for more details.

    You should have received a copy of the GNU General Public License
    along with OpenFOAM.  If not, see <http://www.gnu.org/licenses/>.

\*---------------------------------------------------------------------------*/

#include "turbulentDispersion.H"
#include "addToRunTimeSelectionTable.H"
#include "coupledToConstantDensityFluid.H"
#include "coupledToFluid.H"
#include "standardNormal.H"
#include "wallPolyPatch.H"
#include "maxLagrangianFieldSources.H"
#include "NaNLagrangianFieldSources.H"
#include "internalLagrangianFieldSources.H"

// * * * * * * * * * * * * * * Static Data Members * * * * * * * * * * * * * //

namespace Foam
{
namespace Lagrangian
{
    defineTypeNameAndDebug(turbulentDispersion, 0);
    addToRunTimeSelectionTable
    (
        LagrangianModel,
        turbulentDispersion,
        dictionary
    );
}
}

// * * * * * * * * * * * * * Private Member Functions  * * * * * * * * * * * //

template<class Type>
Foam::tmp<Foam::LagrangianDynamicField<Type>>
Foam::Lagrangian::turbulentDispersion::initialiseTurbField
(
    const word& name,
    const dimensionSet& dims,
    const Type& value
)
{
    return
        tmp<LagrangianDynamicField<Type>>
        (
            new LagrangianDynamicField<Type>
            (
                IOobject
                (
                    name,
                    mesh().time().name(),
                    mesh(),
                    IOobject::READ_IF_PRESENT,
                    IOobject::AUTO_WRITE
                ),
                mesh(),
                dimensioned<Type>(name, dims, value)
            )
        );
}


template<class InjectionFieldSourceType, class Type>
void Foam::Lagrangian::turbulentDispersion::completeTurbField
(
    LagrangianDynamicField<Type>& turbField
)
{
    if (turbField.headerOk()) return;

    LagrangianModels& modelList = cloud().LagrangianModels();

    turbField.sourcesRef().table().transfer
    (
        typename LagrangianDynamicField<Type>::Sources
        (
            turbField,
            modelList.modelTypeFieldSourceTypes
            <
                LagrangianInjection,
                InjectionFieldSourceType,
                LagrangianSource,
                internalLagrangianFieldSource<Type>
            >(),
            IOerrorLocation()
        ).table()
    );
}


// * * * * * * * * * * * * * * * * Constructors  * * * * * * * * * * * * * * //

Foam::Lagrangian::turbulentDispersion::turbulentDispersion
(
    const word& name,
    const LagrangianMesh& mesh,
    const dictionary& modelDict,
    const dictionary& stateDict
)
:
    LagrangianModel(name, mesh),
    cloudLagrangianModel(static_cast<const LagrangianModel&>(*this)),
    dragPtr_(nullptr),
    momentumTransportModel_(mesh.mesh().lookupType<momentumTransportModel>()),
    Cmu75_(pow(dimensionedScalar("Cmu", dimless, modelDict, 0.09), 0.75)),
    maxDiscreteEddies_
    (
        modelDict.lookupOrDefault<label>("maxDiscreteEddies", 32)
    ),
    kc_
    (
        cloud<clouds::carried>().carrierField<scalar>
        (
            clouds::carried::nameToCarrierName
            (
                momentumTransportModel_.k()().name()
            ),
            [&]()
            {
                return momentumTransportModel_.k();
            }
        )
    ),
    epsilonc_
    (
        cloud<clouds::carried>().carrierField<scalar>
        (
            clouds::carried::nameToCarrierName
            (
                momentumTransportModel_.epsilon()().name()
            ),
            [&]()
            {
                return momentumTransportModel_.epsilon();
            }
        )
    ),
    fractionTurb_(initialiseTurbField("fractionTurb", dimless, vGreat)),
    tTurb_(initialiseTurbField("tTurb", dimTime, NaN)),
    Uturb_(initialiseTurbField("Uturb", dimVelocity, vector::uniform(NaN))),
    rndGen_("rndGen", stateDict, name, false),
    avgUturbPtr_(nullptr)
{}


// * * * * * * * * * * * * * * Member Functions  * * * * * * * * * * * * * * //

Foam::wordList Foam::Lagrangian::turbulentDispersion::addSupFields() const
{
    return wordList({cloud().U.name()});
}


bool Foam::Lagrangian::turbulentDispersion::addsSupToField
(
    const word& fieldName,
    const word& eqnFieldName
) const
{
    return
        fieldName == cloud().U.name()
     && (
            eqnFieldName == cloud().U.name()
         || eqnFieldName == cloud<clouds::coupled>().Uc.name()
        );
}


void Foam::Lagrangian::turbulentDispersion::postConstruct()
{
    LagrangianModels& modelList = cloud().LagrangianModels();

    dragPtr_ = nullptr;

    forAll(modelList, i)
    {
        if (!isA<drag>(modelList[i])) continue;

        if (dragPtr_ != nullptr)
        {
            FatalErrorInFunction
                << "Multiple drag models found. Turbulent dispersion "
                << "requires exactly one drag model."
                << exit(FatalError);
        }

        dragPtr_ = &refCast<const drag>(modelList[i]);
    }

    if (dragPtr_ == nullptr)
    {
        FatalErrorInFunction
            << "No drag models found. Turbulent dispersion "
            << "requires exactly one drag model."
            << exit(FatalError);
    }

    completeTurbField<maxLagrangianScalarFieldSource>(fractionTurb_);
    completeTurbField<NaNLagrangianScalarFieldSource>(tTurb_);
    completeTurbField<NaNLagrangianVectorFieldSource>(Uturb_);
}


void Foam::Lagrangian::turbulentDispersion::preAddSup
(
    const LagrangianSubScalarField& deltaT,
    const bool final
)
{
    const LagrangianSubMesh& subMesh = deltaT.mesh();

    // References to the evolving fields
    LagrangianSubScalarSubField& fractionTurb = fractionTurb_.ref(subMesh);
    LagrangianSubScalarSubField& tTurb = tTurb_.ref(subMesh);
    LagrangianSubVectorSubField& Uturb = Uturb_.ref(subMesh);
    const LagrangianSubScalarField& fractionTurb0 = fractionTurb.oldTime();
    const LagrangianSubVectorField& Uturb0 = Uturb.oldTime();

    // Update the eddy time-scale
    const LagrangianSubScalarField magUrel
    (
        mag(cloud().U(subMesh) - cloud<clouds::carried>().Uc(subMesh))
    );
    static const dimensionedScalar rootVSmallEpsilon
    (
        dimEnergy/dimMass/dimTime,
        rootVSmall
    );
    static const dimensionedScalar rootVSmallEpsilonU
    (
        dimEnergy/dimMass/dimTime*dimVelocity,
        rootVSmall
    );
    tTurb =
        max
        (
            kc_(subMesh)/max(epsilonc_(subMesh), rootVSmallEpsilon),
            Cmu75_*kc_(subMesh)*sqrt(kc_(subMesh))
           /max(epsilonc_(subMesh)*magUrel, rootVSmallEpsilonU)
        );

    // Velocity fluctuation magnitude
    LagrangianSubScalarField magUturb(sqrt(2.0/3.0*kc_(subMesh)));

    // Modify particles on the walls to account for the turbulent kinetic
    // energy being constrained to zero. This prevents a turbulent velocity
    // fluctuation from pointing through a wall and causing a particle to
    // vibrate as the drag model and the rebound model fight each other.
    const PackedBoolList patchIsWall
    (
        mesh().mesh().boundaryMesh().findIndices<wallPolyPatch>().toc()
    );
    forAll(subMesh, subi)
    {
        const label i = subi + subMesh.start();

        if (mesh().state(i) < LagrangianState::onPatchZero) continue;

        const label patchi =
            static_cast<label>(mesh().state(i))
          - static_cast<label>(LagrangianState::onPatchZero);

        if (!patchIsWall[patchi]) continue;

        tTurb[subi] = rootVSmall;

        magUturb[subi] = 0;
    }

    // Initialise the average turbulent velocities
    avgUturbPtr_.set
    (
        LagrangianSubVectorField::New
        (
            "avgUturb",
            subMesh,
            dimensionedVector(dimVelocity, Zero)
        ).ptr()
    );

    // Consider each particle in turn
    forAll(subMesh, subi)
    {
        // Create independent generators for each particle. That way if the
        // sequence grows or shrinks across iterations, that effect doesn't
        // propagate and affect other particles' sequences
        randomGenerator rndGen
        (
            rndGen_.sampleAB(labelMin, labelMax),
            false
        );
        distributions::standardNormal stdNormal
        (
            rndGen_.sampleAB(labelMin, labelMax),
            false
        );

        // Generate a random direction with a normally distributed magnitude
        auto rndDir = [&rndGen, &stdNormal]()
        {
            const scalar theta =
                rndGen.scalar01()*constant::mathematical::twoPi;
            const scalar z = 2*rndGen.scalar01() - 1;
            const scalar r = sqrt(1 - sqr(z));
            return stdNormal.sample()*vector(r*cos(theta), r*sin(theta), z);
        };

        // Set up sub-stepping
        const scalar Dt = max(deltaT[subi], rootVSmall);
        scalar dt = 0;

        // Continue/complete the previous eddy
        if (fractionTurb0[subi] < 1)
        {
            dt += tTurb[subi]*(1 - fractionTurb0[subi]);

            avgUturbPtr_()[subi] += min(dt/Dt, 1)*Uturb0[subi];
        }

        // Add new eddies across the time-step
        const scalar nEddies = Dt/tTurb[subi];
        if (nEddies < maxDiscreteEddies_)
        {
            while (dt < Dt)
            {
                // Create a turbulent velocity fluctuation for the new eddy
                Uturb[subi] = rndDir()*magUturb[subi];

                // Add this eddy to the average
                avgUturbPtr_()[subi] +=
                    (min(dt + tTurb[subi], Dt) - dt)/Dt*Uturb[subi];

                // Increment the time
                dt += tTurb[subi];
            }

            // Set the fraction to where we got to in the current eddy
            fractionTurb[subi] = 1 - (dt - Dt)/tTurb[subi];
        }
        else
        {
            // Create a turbulent velocity fluctuation for the new eddy
            Uturb[subi] = rndDir()*magUturb[subi];

            // Add the combined effect of all the eddies to the average
            avgUturbPtr_()[subi] +=
                (Dt - dt)/Dt
               *sqrt(nEddies/3)
               *stdNormal.sample<vector>()
               *magUturb[subi];

            // Set the fraction to indicate that the eddies are finished
            fractionTurb[subi] = 1;
        }
    }

    // If this is not the final iteration then rewind the generator so that the
    // same numbers are generated in the next iteration
    rndGen_.start(!final);
}


void Foam::Lagrangian::turbulentDispersion::addSup
(
    const LagrangianSubScalarField& deltaT,
    const LagrangianSubVectorSubField& U,
    LagrangianEqn<vector>& eqn
) const
{
    assertCloud<clouds::coupledToConstantDensityFluid>();

    eqn.Su += dragPtr_->D(deltaT.mesh())*avgUturbPtr_();
}


void Foam::Lagrangian::turbulentDispersion::addSup
(
    const LagrangianSubScalarField& deltaT,
    const LagrangianSubScalarSubField& vOrM,
    const LagrangianSubVectorSubField& U,
    LagrangianEqn<vector>& eqn
) const
{
    assertCloud
    <
        clouds::coupledToConstantDensityFluid,
        clouds::coupledToFluid
    >();

    eqn.Su += dragPtr_->D(deltaT.mesh())*avgUturbPtr_();
}


void Foam::Lagrangian::turbulentDispersion::postAddSup
(
    const LagrangianSubScalarField& deltaT,
    const bool final
)
{
    avgUturbPtr_.clear();
}


void Foam::Lagrangian::turbulentDispersion::writeProcessorState
(
    Ostream& os
) const
{
    LagrangianModel::writeProcessorState(os);

    writeEntry(os, "rndGen", rndGen_);
}


// ************************************************************************* //
